package tree;

import org.junit.Test;

import java.util.Arrays;

public class KthLargestInBinarySearchTree230 {
    @Test
    public void test() {
        kthLargest(TreeNode.buildTree(Arrays.asList(3,1,4,null,2)), 1);
        kthLargest(TreeNode.buildTree(Arrays.asList(5,3,6,2,4,null,null,1)), 3);
    }

    static class State {
        boolean found;
        int value;
        int count;

        public static State notFoundState(int count) {
            return new State(false, Integer.MIN_VALUE, count);
        }

        public static State foundState(int value) {
            return new State(true, value, -1);
        }

        public State(boolean found, int value, int count) {
            this.found = found;
            this.value = value;
            this.count = count;
        }
    }

    public int kthLargest(TreeNode root, int k) {
        State s = searchKthLargest(root, k);
        if (s.found) {
            System.out.format("found, kth: %d\n", s.value);
            return s.value;
        } else {
            System.out.format("notfound\n");
            return Integer.MIN_VALUE;
        }
    }

    public State searchKthLargest(TreeNode root, int k) {
        if (root == null) {
            return State.notFoundState(0);
        }
        State rightState = searchKthLargest(root.right, k);
        if (rightState.found) {
            return rightState;
        }
        if (k - rightState.count == 1) {
            return State.foundState(root.val);
        }
        State leftState = searchKthLargest(root.left, k - rightState.count - 1);
        if (leftState.found) {
            return leftState;
        }
        return State.notFoundState(leftState.count + rightState.count + 1);
    }
}
